import numpy as np
from scipy.integrate import odeint
from simulate import dissipative_constant, dissipative_force, Hertz_constant, Hertz_force 
from pdb import set_trace

g = 9.8             # gravity (m/s^2)

def Hookean_force(x,k):
    """
    Force of spring following Hooke's law
    x - extension/compression from equilibrium length (m)
    k - spring constant (N/m)
    return force (N)
    """
    F = k * x
    return F

def dashpot_force(x,x_prime,mu):
    """ 
    Dissipative force of the Hookean spring-dashpot model
    x - extension/compression of the Hookean spring (m)
    x_prime - rate of extension/compression of the Hookean spring (m/s)
    mu - proportional constant for disspative force (N s/m)
    return force (N)
    """
    F = mu * x_prime
    return F

def derivatives(y,t,m_s,K,eta,m_c,k,mu):
    """
    function to calculate the derivatives
    y - dependent variable of coupled ODEs -
        [deformation depth, rate change deformation depth, 
         spring compression, rate change spring compression] 
        - [m,m/s,m,m/s]
    t - discrete time array for numerical integration (s)
    m_s - mass of the sphere (kg)
    m_c - mass of the crystal (kg)
    k - Hookean constant (N / m)
    K - Hertz_constant ( kg/(m^0.5*s^2) )
    eta - dissipative constant for Hertz ( kg*m^(1/4)/s )
    mu - dissipative constant for dashpot
    return [2x1] numpy array derivatives
    """
    d0 = y[0]
    d1 = y[1]
    D0 = y[2]
    D1 = y[3]
    try:
        D0_prime = D1
        D1_prime = -1./m_c * (-1 * Hertz_force(d0,K) 
                 + Hookean_force(D0,k) + dashpot_force(D0,D1,mu))
        d0_prime = d1
        d1_prime = -1./m_s * (Hertz_force(d0,K) 
                 + dissipative_force(eta,d0,d1) - m_s * 9.8 
                 + m_s * D1_prime)
        derivative = np.array([d0_prime,d1_prime,D0_prime,D1_prime])
    except (SystemExit, KeyboardInterrupt):
        raise
    return derivative 

def simulate(t=np.linspace(0,10**-3,10**5), m=0.407e-3, R=4.76e-3, 
        h=4.7e-2, Es=3.0e9, Ep=3.0e9, vs=0.33, vp=0.33, COR=0.8,
        mc=4.5e-2, kc=4.0e8, mu=200, cal=6.5e-12):
    """
    Calculate the deformation, deformation rate and force
    at time points given in t for passed parameters.
    t - array of numpy time points (s)
    h - drop height (m)
    R - sphere radius (m)
    Es - sphere modulus (Pa)
    Ep - plane modulus (Pa)
    vs - sphere Poisson's ratio
    vp - plane Poisson's ratio
    m - sphere mass (kg)
    COR - coefficient of restitution
    mc : mass of the crystal/compact/load cell (kg)
    kc : stiffness of the spring (N/m)
    mu : damping factor of the crystal (N s/m)
    cal : load cell calibration (m/mV)
    """
    K = Hertz_constant(Es,Ep,vs,vp,R)
    eta = dissipative_constant(m,K,COR)
    vo = g*(2*h/g)**0.5
    y_init = [0,1.0*vo,0,0.0*vo]
    y = odeint(derivatives,y_init,t,(m,K,eta,mc,kc,mu))
    mask = ~np.isnan(y[:,0])
    t = t[mask]
    y = y[mask]
    pts = np.sum(mask)
    step = max(1, pts / 1000)
    t = np.array([t[i] for i in range(0,pts,step)])
    y = np.array([y[i,:] for i in range(0,pts,step)])
    f = Hertz_force(y[:,0],K) 
    mask = ~np.isnan(f)
    t = t[mask]
    y = y[mask]
    f = f[mask]
    dynamics = np.zeros((t.shape[0]), dtype=[('time','<f4'), 
        ('deformation','<f4'), ('deformation rate', '<f4'), 
        ('deflection','<f4'), ('deflection rate', '<f4'), 
        ('force', '<f4'), ('voltage','<f4')])
    dynamics['time'] = t 
    dynamics['deformation'] = y[:,0] 
    dynamics['deformation rate'] = y[:,1] 
    dynamics['deflection'] = y[:,2] 
    dynamics['deflection rate'] = y[:,3] 
    dynamics['force'] = f 
    dynamics['voltage'] = 1./cal * dynamics['deflection']
    return dynamics
